#!/usr/bin/python

# Hmmtop object runs hmmtop on instantiation returns an object of parsed
# results from hmmtop run
#
# hmmtop constructor:
#   header - fasta header line
#   seq - fasta sequence
#
# hmmtop instance vars:
#   success    - Flag to see if run was successful
#   fasta_seq  - [header, seq] of input sequence
#   output     - raw hmmtop output
#   num_tms    - number of tms
#   tms_ranges - indicies of tms sequences
#   tms_seq    - tms sequences
#   is_tms     - array of bool indicating true where there is tms

import subprocess
import tempfile
import os

HMMTOP = '/Users/kevin/saier/hmmtop_2_1/'
HMMTOP_BIN = HMMTOP + 'hmmtop'
HMMTOP_ARCH = HMMTOP + 'hmmtop.arch'
HMMTOP_PSV = HMMTOP + 'hmmtop.psv'

os.environ['HMMTOP_ARCH'] = HMMTOP_ARCH
os.environ['HMMTOP_PSV'] = HMMTOP_PSV

class hmmtop:

  def __init__(self, header, seq):
    self.success = False            # flag to mark run success
    if header[0] != '>':
      header = '>' + header
    self.header = header
    self.seq = seq

    # Genereate a temporary file for Hmmtop input
    temp_input = tempfile.NamedTemporaryFile()
    temp_input.write(header + "\n" + seq)
    temp_input.seek(0)
    hmmtop_cmd = HMMTOP_BIN + " -if=%s -pl -sf=FAS" % temp_input.name
    # System call to Hmmtop, output contains stdout and stderr
    hmmtop_output = subprocess.Popen(hmmtop_cmd, stdout=subprocess.PIPE, \
        stderr=subprocess.PIPE, shell=True).communicate()
    temp_input.close()

    # Save hmmtop output in obj in case needed
    self.output = hmmtop_output

    # Make sure hmmtop ran correctly
    if hmmtop_output[0] == "":
      print "Error occured when running Hmmtop, check your input."
      return

    hmmtop_output = hmmtop_output[0].split('\n')[0]   # get the first line
    if " IN " in hmmtop_output:
      hmmtop_output = hmmtop_output[hmmtop_output.rindex(" IN ")+4:].split()
    elif " OUT " in hmmtop_output:
      hmmtop_output = hmmtop_output[hmmtop_output.rindex(" OUT ")+5:].split()
    else:
      return

    self.num_tms = int(hmmtop_output[0])      # save the number of tms
    hmmtop_output = hmmtop_output[1:]         # reassign to tms ranges

    # loop through tms region ranges and assign them to a new array
    self.tms_ranges = []
    for i in range(0, len(hmmtop_output), 2):
      self.tms_ranges.append((int(hmmtop_output[i])-1, \
          int(hmmtop_output[i+1])))

    # get the residues at these tms
    self.tms_seq = []
    for i in self.tms_ranges:
      self.tms_seq.append(seq[i[0]:i[1]])

    # create an array with true where there are tms regions
    self.is_tms = [False]*len(seq)
    for i in self.tms_ranges:
      for j in range(i[0], i[1]):
        self.is_tms[j] = True

    self.success = True            # if gets here, everything is success

    # pass in a subsequence of aligned sequence with gaps
    # returns an array of true/false/- where gaps are inserted into is_tms subset
    #    if gapped seq not in original passed in seq, return false
    # params: string sequence with gaps
  def get_is_tms(self, seq):

    # get sequence with no gaps
    seq_nogap = seq.replace('-', '')
    # check if passed in seq is in the hmmtop obj
    if not seq_nogap in self.seq:
      return False

    # get sequence start and end
    seq_beg = self.seq.index(seq_nogap)
    seq_end = seq_beg + len(seq_nogap)
    is_tms_subseq = self.is_tms[seq_beg:seq_end]

    # insert '-' into is_tms where is not true or false
    for i in range(len(seq)):
      if seq[i] == '-':
        is_tms_subseq = is_tms_subseq[:i] + ['-'] + is_tms_subseq[i:]

    return is_tms_subseq 


    
